package org.fujay.langchain4jserver.infrastructure.persistence.assembler;

import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.ObjectMapper;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.fujay.langchain4jserver.domain.conversation.entity.MessageEntity;
import org.fujay.langchain4jserver.domain.conversation.vo.MessageId;
import org.fujay.langchain4jserver.domain.conversation.vo.MessageType;
import org.fujay.langchain4jserver.domain.conversation.vo.TokenUsage;
import org.fujay.langchain4jserver.infrastructure.persistence.mybatis.entity.MessagePo;
import org.springframework.stereotype.Component;

import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

@Slf4j
@Component
@RequiredArgsConstructor
public class MessageAssembler {

    private final ObjectMapper objectMapper;

    public MessageEntity toDomainEntity(MessagePo po) {
        if (po == null) {
            return null;
        }
        MessageId id = new MessageId(po.getId());
        MessageType type = MessageType.valueOf(po.getMessageType().toUpperCase());

        TokenUsage usage = null;
        if (StringUtils.isNotBlank(po.getTokenUsageJson())) {
            try {
                usage = objectMapper.readValue(po.getTokenUsageJson(), TokenUsage.class);
            } catch (JsonProcessingException e) {
                log.error(
                        "Failed to deserialize TokenUsage JSON for message {}: {}", id.value(), e.getMessage());
            }
        }

        Map<String, Object> metadata = null;
        if (StringUtils.isNotBlank(po.getMetadataJson())) {
            try {
                metadata =
                        objectMapper.readValue(
                                po.getMetadataJson(), new TypeReference<Map<String, Object>>() {
                                });
            } catch (JsonProcessingException e) {
                log.error(
                        "Failed to deserialize Metadata JSON for message {}: {}", id.value(), e.getMessage());
            }
        }

        List<String> imageUrls = null;
        if (StringUtils.isNotBlank(po.getImageUrlsJson())) {
            try {
                imageUrls =
                        objectMapper.readValue(po.getImageUrlsJson(), new TypeReference<List<String>>() {
                        });
            } catch (JsonProcessingException e) {
                log.error(
                        "Failed to deserialize ImageUrls JSON for message {}: {}", id.value(), e.getMessage());
            }
        }

        return MessageEntity.load(
                id, type, po.getContent(), po.getTimestamp(), usage, metadata, imageUrls);
    }

    public List<MessageEntity> toDomainEntities(List<MessagePo> pos) {
        if (pos == null) {
            return Collections.emptyList();
        }
        return pos.stream().map(this::toDomainEntity).collect(Collectors.toList());
    }

    public MessagePo toPO(MessageEntity entity) {
        if (entity == null) {
            return null;
        }
        MessagePo po = new MessagePo();
        po.setId(entity.getId() != null ? entity.getId().value() : null);
        // conversationId is set by the caller (Repository) when saving
        po.setMessageType(entity.getType() != null ? entity.getType().name() : null);
        po.setContent(entity.getContent());
        po.setTimestamp(entity.getTimestamp());

        try {
            if (entity.getTokenUsage() != null) {
                po.setTokenUsageJson(objectMapper.writeValueAsString(entity.getTokenUsage()));
            }
            if (entity.getMetadata() != null && !entity.getMetadata().isEmpty()) {
                po.setMetadataJson(objectMapper.writeValueAsString(entity.getMetadata()));
            }
            if (entity.getImageUrls() != null && !entity.getImageUrls().isEmpty()) {
                po.setImageUrlsJson(objectMapper.writeValueAsString(entity.getImageUrls()));
            }
        } catch (JsonProcessingException e) {
            log.error(
                    "Failed to serialize JSON for message {}: {}",
                    entity.getId() != null ? entity.getId().value() : "<new>",
                    e.getMessage());
        }
        return po;
    }

    public List<MessagePo> toPOs(List<MessageEntity> entities) {
        if (entities == null) {
            return Collections.emptyList();
        }
        return entities.stream().map(this::toPO).collect(Collectors.toList());
    }
}
